#ifndef ENGINE_H
#define ENGINE_H

#include <NvInfer.h>

#include <atomic>
#include <memory>
#include <string>

#include "nvbuffers.h"
/* #include "config.h" */

namespace weilin {

class Engine {
 public:
  Engine() : isAvaliable(false) {}

  bool build_engine(const std::string& file);
  bool build_engine(const std::shared_ptr<nvinfer1::IHostMemory>& model);
  bool build_engine(const char* data, const std::size_t size);

  void perpare();

  void execute();

  void execute_async();

  std::shared_ptr<nvinfer1::ICudaEngine> get_engine() { return mEngine; }

  std::shared_ptr<TensorBuffer>& get_buffer();

  bool is_avaliable() const { return isAvaliable; }

 private:
  bool isAvaliable;
  std::shared_ptr<nvinfer1::ICudaEngine> mEngine;
  std::unique_ptr<nvinfer1::IExecutionContext> mContext;
  std::shared_ptr<TensorBuffer> mBuffers;
};

}  // namespace weilin

#endif  // !ENGINE_H
